import tkinter as tk
from tkinter import filedialog, ttk, messagebox
from PIL import Image, ImageTk
import cv2
import json
import os
import numpy as np

class PhotoEditorApp:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("智能证件照编辑器")
        
        # 初始化界面组件
        self.create_widgets()
        self.load_config()
        
    def create_widgets(self):
        # 文件选择区域
        self.btn_select = ttk.Button(self.root, text="选择照片", command=self.select_image)
        self.btn_select.pack(pady=10)
        
        # 尺寸选择
        self.size_var = tk.StringVar()
        self.size_combobox = ttk.Combobox(self.root, textvariable=self.size_var, state='readonly')
        self.size_combobox.pack(pady=5)
        
        # 背景色选择
        self.bg_color_var = tk.StringVar()
        self.color_combobox = ttk.Combobox(self.root, textvariable=self.bg_color_var, state='readonly')
        self.color_combobox.pack(pady=5)
        
        # 预览区域
        self.preview_label = tk.Label(self.root)
        self.preview_label.pack(pady=10)
        
        # 处理按钮
        self.btn_process = ttk.Button(self.root, text="生成证件照", command=self.save_image)
        self.btn_process.pack(pady=10)

    def load_config(self):
        # 加载配置文件
        config_path = os.path.join(os.path.dirname(__file__), 'config.json')
        with open(config_path, encoding='utf-8') as f:
            config = json.load(f)
            
        self.config = config  # 保存配置到实例变量
        self.size_combobox['values'] = list(config['sizes'].keys())
        self.color_combobox['values'] = list(config['background_colors'].keys())
        self.size_var.set('一寸')
        self.bg_color_var.set('白色')

    def select_image(self):
        file_path = filedialog.askopenfilename(filetypes=[("图片文件", "*.jpg *.png")])
        if not file_path:
            return
        try:
            # 确保路径编码正确
            file_path = file_path.encode('utf-8').decode('utf-8')
            self.current_image_path = file_path
        except UnicodeError:
            messagebox.showerror("错误", "无法正确读取文件路径，请检查路径中的特殊字符")
            return

        # 实时更新预览
        self.size_combobox.bind('<<ComboboxSelected>>', lambda event: self.process_image(self.current_image_path))
        self.color_combobox.bind('<<ComboboxSelected>>', lambda event: self.process_image(self.current_image_path))
        self.process_image(self.current_image_path)

    def resize_photo(self, original_img, size_config):
        """智能调整证件照尺寸并保护头部空间"""
        # 计算原始宽高比
        h, w = original_img.shape[:2]
        target_ratio = size_config['width'] / size_config['height']
        
        # 动态调整目标尺寸
        if w/h > target_ratio:
            new_h = size_config['height']
            new_w = int(new_h * w/h)
        else:
            new_w = size_config['width']
            new_h = int(new_w * h/w)
        
        # 高质量缩放
        resized = cv2.resize(original_img, (new_w, new_h), 
                           interpolation=cv2.INTER_LANCZOS4)
        
        # 智能裁剪保护头部空间
        y_start = max(0, (new_h - size_config['height']) // 4)
        y_end = y_start + size_config['height']
        x_center = (new_w - size_config['width']) // 2
        return resized[y_start:y_end, x_center:x_center+size_config['width']]

    def detect_background(self, image):
        """使用HSV色彩空间检测背景区域并优化边缘处理"""
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # 定义白色和蓝色的HSV范围
        lower_white = np.array([0, 0, 200])
        upper_white = np.array([180, 30, 255])
        lower_blue = np.array([90, 60, 60])
        upper_blue = np.array([130, 255, 255])
        
        # 创建遮罩
        mask_white = cv2.inRange(hsv, lower_white, upper_white)
        mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
        combined_mask = cv2.bitwise_or(mask_white, mask_blue)
        
        # 形态学操作优化遮罩
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
        refined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel)
        refined_mask = cv2.morphologyEx(refined_mask, cv2.MORPH_OPEN, kernel)
        
        # 边缘检测优化，增加高斯模糊和膨胀操作
        blurred = cv2.GaussianBlur(refined_mask, (5, 5), 0)
        edges = cv2.Canny(blurred, 50, 150)
        kernel_dilate = np.ones((3, 3), np.uint8)
        dilated_edges = cv2.dilate(edges, kernel_dilate, iterations=1)
        return cv2.bitwise_or(refined_mask, dilated_edges)

    def replace_background(self, image, mask, bg_type, color):
        """替换背景支持纯色/渐变效果"""
        if bg_type == '纯色':
            background = np.full_like(image, color)
        elif bg_type == '渐变':
            # 创建垂直渐变背景
            grad = np.linspace(color[0]-30, color[0], image.shape[0])
            background = np.zeros_like(image)
            background[:,:,0] = grad[:,np.newaxis]
            background[:,:,1] = color[1]
            background[:,:,2] = color[2]
        
        # 使用遮罩混合图像
        return cv2.bitwise_and(image, image, mask=cv2.bitwise_not(mask)) + \
               cv2.bitwise_and(background, background, mask=mask)

    def process_image(self, image_path=None):
        if not image_path:
            return
        
        try:
            # 读取原始图片
            image_path = os.path.normpath(image_path)
            original_img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
            if original_img is None:
                messagebox.showerror("错误", "无法读取图片文件")
                return
            
            # 获取配置参数
            size_config = self.config['sizes'][self.size_var.get()]
            bg_config = self.config['background_colors'][self.bg_color_var.get()]
            
            # 尺寸转换
            resized_img = self.resize_photo(original_img, size_config)
            
            # 背景检测
            mask = self.detect_background(resized_img)
            
            # 背景替换
            self.result = self.replace_background(resized_img, mask, 
                                                bg_config['type'], bg_config['color'])
            
            # 显示预览
            preview_img = cv2.cvtColor(self.result, cv2.COLOR_BGR2RGB)
            preview_img = Image.fromarray(preview_img)
            preview_img.thumbnail((400, 500))
            photo = ImageTk.PhotoImage(preview_img)
            
            self.preview_label.configure(image=photo)
            self.preview_label.image = photo
            
        except Exception as e:
            if not hasattr(self, 'result'):
                # 使用默认黑色背景
                self.result = np.zeros((300, 300, 3), dtype=np.uint8)
            
            messagebox.showerror("错误", f"处理失败: {str(e)}")
            
            # 显示错误状态预览
            preview_img = cv2.cvtColor(self.result, cv2.COLOR_BGR2RGB)
            preview_img = Image.fromarray(preview_img)
            preview_img.thumbnail((400, 500))
            photo = ImageTk.PhotoImage(preview_img)
            
            self.preview_label.configure(image=photo)
            self.preview_label.image = photo
            
            # 重置配置加载
            self.load_config()
            
        except Exception as e:
            print(f"处理图片时发生错误: {str(e)}")
            messagebox.showerror("错误", f"处理图片时发生错误: {str(e)}")
            return

    def save_image(self):
        if not hasattr(self, 'current_image_path'):
            return

        save_name = os.path.splitext(os.path.basename(self.current_image_path))[0]
        save_ext = '.jpg'
        # 修改为使用 current_image_path 的目录
        save_path = os.path.join(os.path.dirname(self.current_image_path), f'{save_name}_edited{save_ext}')

        # 保存图片
        preview_img = cv2.cvtColor(self.result, cv2.COLOR_BGR2RGB)
        preview_img = Image.fromarray(preview_img)
        preview_img.save(save_path, format='JPEG', quality=95)
        messagebox.showinfo("成功", f"证件照已保存至：{save_path}")

if __name__ == "__main__":
    app = PhotoEditorApp()
    app.root.mainloop()